import cv2
import numpy as np
import time
from picamera2 import Picamera2

class EnhancedObstacleVisualizer:
    def __init__(self):
        # 颜色定义
        self.COLORS = {
            'edge': (0, 255, 255),    # 黄色 - 边缘检测
            'motion': (255, 0, 0),    # 蓝色 - 运动检测
            'depth': (0, 0, 255),     # 红色 - 深度检测
            'safe_area': (0, 255, 0), # 绿色 - 安全区域
            'warning': (0, 165, 255), # 橙色 - 警告区域
            'path': (0, 255, 0),      # 绿色 - 安全路径
            'distance_line': (200, 200, 200)  # 灰色 - 距离参考线
        }
        
        # 初始化背景减法器
        self.fgbg = cv2.createBackgroundSubtractorMOG2(history=50, varThreshold=16, detectShadows=False)
        
        # 避障参数
        self.obstacle_threshold = 0.25
        self.safe_distance_ratio = 0.6
        self.left_safe = False
        self.right_safe = False
    
    def visualize_and_detect(self, frame):
        h, w = frame.shape[:2]
        visualization = frame.copy()
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        
        # 1. 绘制距离参考线
        for i in range(1, 4):
            y_pos = int(h * (1 - i*0.2))  # 20%, 40%, 60%高度
            cv2.line(visualization, (0, y_pos), (w, y_pos), self.COLORS['distance_line'], 1)
            cv2.putText(visualization, f"{i*0.5}m", (10, y_pos-5), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.4, self.COLORS['distance_line'], 1)
        
        # 2. 边缘检测可视化
        edges = cv2.Canny(gray, 50, 150)
        roi_edges = edges[int(h*0.6):, :]
        edge_contours, _ = cv2.findContours(roi_edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for cnt in edge_contours:
            if cv2.contourArea(cnt) > 100:
                cnt[:, :, 1] += int(h*0.6)
                cv2.drawContours(visualization, [cnt], -1, self.COLORS['edge'], 2)
        
        # 3. 运动检测可视化
        fgmask = self.fgbg.apply(frame)
        fgmask = cv2.erode(fgmask, None, iterations=1)
        fgmask = cv2.dilate(fgmask, None, iterations=2)
        motion_contours, _ = cv2.findContours(fgmask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for cnt in motion_contours:
            if cv2.contourArea(cnt) > 500:
                x, y, w, h_rect = cv2.boundingRect(cnt)
                cv2.rectangle(visualization, (x, y), (x+w, y+h_rect), self.COLORS['motion'], 2)
        
        # 4. 深度检测可视化
        _, thresh = cv2.threshold(gray[int(h*0.5):, :], 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
        depth_contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for cnt in depth_contours:
            if cv2.contourArea(cnt) > 500:
                cnt[:, :, 1] += int(h*0.5)
                x, y, w, h_rect = cv2.boundingRect(cnt)
                cv2.drawContours(visualization, [cnt], -1, self.COLORS['depth'], 2)
                # 添加距离估计
                distance = int(1 / (h_rect/h) * 0.5)
                cv2.putText(visualization, f"{distance}m", (x, y-5), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
        
        # 5. 区域安全分析
        left_region = gray[int(h*0.6):, :w//2]
        right_region = gray[int(h*0.6):, w//2:]
        self.left_safe = self.analyze_region_safety(left_region)
        self.right_safe = self.analyze_region_safety(right_region)
        
        # 6. 绘制区域划分
        cv2.rectangle(visualization, (0, int(h*0.6)), (w, h), self.COLORS['warning'], 2)
        cv2.rectangle(visualization, (0, 0), (w, int(h*0.6)), self.COLORS['safe_area'], 2)
        cv2.line(visualization, (w//2, 0), (w//2, h), (255, 255, 255), 1)
        
        # 7. 绘制动态避障路径
        self.draw_navigation_path(visualization, h, w)
        
        # 8. 添加安全状态文本
        cv2.putText(visualization, f"Left: {'SAFE' if self.left_safe else 'DANGER'}", 
                   (10, h-30), cv2.FONT_HERSHEY_SIMPLEX, 0.5, 
                   self.COLORS['safe_area'] if self.left_safe else self.COLORS['warning'], 2)
        cv2.putText(visualization, f"Right: {'SAFE' if self.right_safe else 'DANGER'}", 
                   (w//2+10, h-30), cv2.FONT_HERSHEY_SIMPLEX, 0.5, 
                   self.COLORS['safe_area'] if self.right_safe else self.COLORS['warning'], 2)
        
        return visualization
    
    def analyze_region_safety(self, region):
        edges = cv2.Canny(region, 50, 150)
        edge_density = np.sum(edges > 0) / (edges.size + 1e-6)
        return edge_density < 0.1
    
    def draw_navigation_path(self, img, h, w):
        if self.left_safe and not self.right_safe:
            cv2.arrowedLine(img, (w//2, h), (w//4, h//2), self.COLORS['path'], 2, tipLength=0.3)
        elif self.right_safe and not self.left_safe:
            cv2.arrowedLine(img, (w//2, h), (3*w//4, h//2), self.COLORS['path'], 2, tipLength=0.3)
        elif self.left_safe and self.right_safe:
            cv2.arrowedLine(img, (w//2, h), (w//2, h//2), self.COLORS['path'], 2, tipLength=0.3)
        else:
            cv2.putText(img, "STOP!", (w//2-30, h//2), 
                       cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

class EnhancedVisionCarController:
    def __init__(self):
        self.visualizer = EnhancedObstacleVisualizer()

    def calculate_edge_density(self, region):
        """
        计算图像区域的边缘密度（边缘像素占比）
        参数:
            region: 输入图像区域 (灰度或BGR)
        返回:
            float: 边缘密度值 (0-1之间)
        """
        # 如果是彩色图像，先转换为灰度
        if len(region.shape) == 3:
            gray = cv2.cvtColor(region, cv2.COLOR_BGR2GRAY)
        else:
            gray = region
    
        # 边缘检测 (Canny算法)
        edges = cv2.Canny(gray, 50, 150)
    
        # 计算边缘像素占比
        edge_pixels = np.sum(edges > 0)
        total_pixels = edges.size
        density = edge_pixels / (total_pixels + 1e-6)  # 加极小值避免除零
        return density    
    
    def process_frame(self, frame):
        # 可视化并检测障碍物
        visualized = self.visualizer.visualize_and_detect(frame)
        # 控制决策
        if not self.visualizer.left_safe and not self.visualizer.right_safe:
            # 完全阻塞，原地左转
            action = "Full Block - Turning Left"
            time.sleep(0.5)
        elif not self.visualizer.left_safe:
            # 右侧阻塞，左转
            action = "Right Block - Turning Left"
            time.sleep(0.3)
        elif not self.visualizer.right_safe:
            # 左侧阻塞，右转
            action = "Left Block - Turning Right"
            time.sleep(0.3)
        else:
            # 无障碍时，检查安全区域是否居中
            h, w = frame.shape[:2]
            left_region = frame[int(h*0.6):, :w//2]
            right_region = frame[int(h*0.6):, w//2:]
            left_density = self.calculate_edge_density(left_region)
            right_density = self.calculate_edge_density(right_region)
            cv2.putText(visualized, f"L-Den: {left_density:.2f}", (10, 200),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
            cv2.putText(visualized, f"R-Den: {right_density:.2f}", (w//2, 200),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
            # 根据两侧密度差调整方向
            if abs(left_density - right_density) > 0.02:  # 明显不平衡
                if left_density < right_density:
                    action = "Turning Left"
                else:
                    action = "Turning Right"
            else:
               # 无障碍，前进
                action = "Clear Path - Moving Forward"
        
        cv2.putText(visualized, f"Action: {action}", (10, 30), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
        
        return visualized
   
  
if __name__ == "__main__":
    car = EnhancedVisionCarController()
    picam2 = Picamera2()
    picam2.configure(picam2.create_preview_configuration(main={"format": 'XRGB8888', "size": (640,480)}))
    picam2.start()
    frame =picam2.capture_array()
    cv2.imshow("Frame",  car.process_frame(frame))
